{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HMinCaml.Beta
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Beta reduction.
--
--------------------------------------------------------------------------------

module HMinCaml.Beta
  (

    beta

  ) where

import qualified HMinCaml.HMonad as H
import HMinCaml.Id
import HMinCaml.KNormal


import Control.Applicative
import Control.Monad
import qualified Data.Map as M



--------------------------------------------------------------------------------
-- Beta reduction


-- Same monad as Alpha - could share code...
--
type Env = M.Map IdS IdS

newtype Mon a = Mon { getMon :: Env -> H.HMonad a }


instance Functor Mon where
  fmap f ma = Mon $ \r -> fmap f $ getMon ma r

instance Applicative Mon where
  pure a    = Mon $ \_ -> pure a
  af <*> av = Mon $ \r -> getMon af r <*> getMon av r

instance Monad Mon where
  return    = pure
  ma >>= k  = Mon $ \r -> getMon ma r >>= \a -> getMon (k a) r

instance MonadPlus Mon where
  mzero         = Mon $ \_ -> mzero
  ma `mplus` mb = Mon $ \r -> getMon ma r `mplus` getMon mb r

-- Not modular - TODO should not have a run function...
--
runMon :: Mon a -> H.HMonad a
runMon ma = getMon ma M.empty

local :: (Env -> Env) -> Mon a -> Mon a
local f ma = Mon $ \r -> getMon ma (f r)


localBinding :: IdS -> IdS -> Mon a -> Mon a
localBinding v v1 ma = local (M.insert v v1) ma

find :: IdS -> Mon IdS
find x = Mon $ \r -> maybe (return x) return $ M.lookup x r
    

beta :: Expr -> H.HMonad Expr
beta expr = runMon $ bc expr


bc :: Expr -> Mon Expr
bc (Unit)                 = return Unit
bc (Int i)                = return $ Int i
bc (Float d)              = return $ Float d
bc (Neg x)                = Neg  <$> find x
bc (Add x y)              = Add  <$> find x <*> find y
bc (Sub x y)              = Sub  <$> find x <*> find y
bc (FNeg x)               = FNeg <$> find x
bc (FAdd x y)             = FAdd <$> find x <*> find y
bc (FSub x y)             = FSub <$> find x <*> find y
bc (FMul x y)             = FMul <$> find x <*> find y
bc (FDiv x y)             = FDiv <$> find x <*> find y
bc (IfEq x y e1 e2)       = IfEq <$> find x <*> find y <*> bc e1 <*> bc e2
bc (IfLE x y e1 e2)       = IfLE <$> find x <*> find y <*> bc e1 <*> bc e2
bc (Let (x,t) e1 e2)      = bc e1 >>= \ans -> 
    case ans of 
       Var y -> localBinding x y $ bc e2
       ea    -> Let (x,t) ea <$> bc e2

bc (Var x)                = Var <$> find x

bc (LetRec (Fundef (x,t) yts e1) e2) = do
    (\ea eb -> LetRec (Fundef (x,t) yts ea) eb) <$> bc e1 <*> bc e2

bc (App x ys)               = App <$> find x <*> mapM find ys
bc (Tuple xs)               = Tuple <$> mapM find xs
bc (LetTuple xts y e)       = LetTuple xts <$> find y <*> bc e
bc (Get x y)                = Get <$> find x <*> find y
bc (Put x y z)              = Put <$> find x <*> find y <*> find z
bc (ExtArray x)             = ExtArray <$> find x
bc (ExtFunApp x ys)         = ExtFunApp <$> find x <*> mapM find ys

